package onnxruntime

import (
	"errors"
)

// #cgo CFLAGS: -O2 -g
//
// #include "onnxruntime_wrapper.h"
import "C"

// Holds options required when enabling the CUDA backend for a session. This
// struct wraps C onnxruntime types; users must create instances of this using
// the NewCUDAProviderOptions() function. So, to enable CUDA for a session,
// follow these steps:
//
//  1. Call NewSessionOptions() to create a SessionOptions struct.
//  2. Call NewCUDAProviderOptions() to obtain a CUDAProviderOptions struct.
//  3. Call the CUDAProviderOptions struct's Update(...) function to pass a
//     list of settings to CUDA. (See the comment on the Update() function.)
//  4. Pass the CUDA options struct pointer to the
//     SessionOptions.AppendExecutionProviderCUDA(...) function.
//  5. Call the Destroy() function on the CUDA provider options.
//  6. Call NewAdvancedSession(...), passing the SessionOptions struct to it.
//  7. Call the Destroy() function on the SessionOptions struct.
//
// Admittedly, this is a bit of a mess, but that's how it's handled by the C
// API internally. (The onnxruntime python API hides a bunch of this complexity
// using getter and setter functions, for which Go does not have a terse
// equivalent.)

// Initializes and returns a CUDAProviderOptions struct, used when enabling
// CUDA in a SessionOptions instance. (i.e., a CUDAProviderOptions must be
// configured, then passed to SessionOptions.AppendExecutionProviderCUDA.)
// The caller must call the Destroy() function on the returned struct when it's
// no longer needed.
func NewCUDAProviderOptions() (*CUDAProviderOptions, error) {
	if !IsInitialized() {
		return nil, ErrorNotInitialized
	}
	var o *C.OrtCUDAProviderOptionsV2
	status := C.CreateCUDAProviderOptions(&o)
	if status != nil {
		return nil, statusToError(status)
	}
	return &CUDAProviderOptions{
		o: o,
	}, nil
}

type CUDAProviderOptions struct {
	o *C.OrtCUDAProviderOptionsV2
}

// Wraps the call to the UpdateCUDAProviderOptions in the onnxruntime C API.
// Requires a map of string keys to values for configuring the CUDA backend.
// For example, set the key "device_id" to "1" to use GPU 1 rather than 0.
//
// The onnxruntime headers refer users to
// https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#configuration-options
// for a full list of available keys and values.
func (o *CUDAProviderOptions) Update(options map[string]string) error {
	if len(options) == 0 {
		return nil
	}
	keys, values := mapToCStrings(options)
	defer freeCStrings(keys)
	defer freeCStrings(values)
	status := C.UpdateCUDAProviderOptions(o.o, &(keys[0]), &(values[0]),
		C.int(len(options)))
	if status != nil {
		return statusToError(status)
	}
	return nil
}

// Must be called when the CUDAProviderOptions struct is no longer needed;
// frees internal C-allocated state. Note that the CUDAProviderOptions struct
// can be destroyed as soon as options.AppendExecutionProviderCUDA has been
// called.
func (o *CUDAProviderOptions) Destroy() error {
	if o.o == nil {
		return errors.New("the CUDAProviderOptions are not initialized")
	}
	C.ReleaseCUDAProviderOptions(o.o)
	o.o = nil
	return nil
}

// Initializes and returns a TensorRTProviderOptions struct, used when enabling
// the TensorRT backend. The caller must call the Destroy() function on the
// returned struct when it's no longer needed.
func NewTensorRTProviderOptions() (*TensorRTProviderOptions, error) {
	if !IsInitialized() {
		return nil, ErrorNotInitialized
	}
	var o *C.OrtTensorRTProviderOptionsV2
	status := C.CreateTensorRTProviderOptions(&o)
	if status != nil {
		return nil, statusToError(status)
	}
	return &TensorRTProviderOptions{
		o: o,
	}, nil
}

// Like the CUDAProviderOptions struct, but used for configuring TensorRT
// options. Instances of this struct must be initialized using
// NewTensorRTProviderOptions() and cleaned up by calling their Destroy()
// function when they are no longer needed.
type TensorRTProviderOptions struct {
	o *C.OrtTensorRTProviderOptionsV2
}

// Wraps the call to the UpdateTensorRTProviderOptions in the C API. Requires
// a map of string keys to values.
//
// The onnxruntime headers refer users to
// https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html#cc
// for the list of available keys and values.
func (o *TensorRTProviderOptions) Update(options map[string]string) error {
	if len(options) == 0 {
		return nil
	}
	keys, values := mapToCStrings(options)
	defer freeCStrings(keys)
	defer freeCStrings(values)
	status := C.UpdateTensorRTProviderOptions(o.o, &(keys[0]), &(values[0]),
		C.int(len(options)))
	if status != nil {
		return statusToError(status)
	}
	return nil
}

// Must be called when the TensorRTProviderOptions are no longer needed, in
// order to free internal state. The struct is not needed as soon as you have
// passed it to the AppendExecutionProviderTensorRT function.
func (o *TensorRTProviderOptions) Destroy() error {
	if o.o == nil {
		return errors.New("the TensorRTProviderOptions are not initialized")
	}
	C.ReleaseTensorRTProviderOptions(o.o)
	o.o = nil
	return nil
}
